{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zdbPo0MbQ4R7"
      },
      "source": [
        "<img src=\"https://github.com/stanfordnlp/dspy/blob/main/docs/images/DSPy8.png?raw=1\" alt=\"DSPy7 Image\" height=\"150\"/>\n",
        "\n",
        "### **SolveGSM8k**: Solving grade school math problems using DSPy\n",
        "\n",
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/stanfordnlp/dspy/blob/main/examples/math/gsm8k/gsm8k_assertions.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Aq1HjnZdV0e"
      },
      "source": [
        "This notebook builds upon the foundational concepts of the DSPy framework. DSPy overs a novel programming-centric approach to utilizing language and retrieval models. It offers a unique blend of prompting, reasoning, fine-tuning, and tool augmentation, all encapsulated under a minimalistic Python syntax.\n",
        "\n",
        "We will focus on three parts:\n",
        "\n",
        "\n",
        "1.   Define a DSPy program and evaluate its performance.\n",
        "3.   Constrain DSPy program's behavior with runtime DSPy assertions and suggestions.\n",
        "2.   Optimize the DSPy program with in-context learning and prompt tuning.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NckBcbZN9p6y"
      },
      "outputs": [],
      "source": [
        "# Set up your API key for OpenAI\n",
        "import os\n",
        "os.environ[\"OPENAI_API_KEY\"]=\"Paste your key here\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ghG4tCe3e-kQ"
      },
      "source": [
        "### Step -1. **Installing Cache and DSPy** (Run the collapsed cells)\n",
        "\n",
        "The cells are collapsed by default. Running the following cells will set up cache and install all dependencies and DSPy.\n",
        "\n",
        "The first cell ensures all the following LM calls in this notebook will be using the cached OpenAI's API result. Removing this step might sigificantly increase the running time of all the following DSPy programs depending on your OpenAI account setup.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gXLDcxrLPgBO"
      },
      "outputs": [],
      "source": [
        "# This cell sets up the cache (pre-computed OpenAI call results).\n",
        "!rm -r gsm8k_cache || true\n",
        "!rm -r dspy || true\n",
        "!git clone https://github.com/Shangyint/gsm8k_cache.git\n",
        "\n",
        "import os\n",
        "repo_clone_path = '/content/gsm8k_cache'\n",
        "\n",
        "# Check if '/content' is writable\n",
        "if not os.access('/content', os.W_OK):\n",
        "    # If '/content' is not writable, choose an alternative directory\n",
        "    # Example: using a directory relative to the current working directory\n",
        "    repo_clone_path = os.path.join(os.getcwd(), 'gsm8k_cache')\n",
        "\n",
        "# Set up the cache for this notebook\n",
        "os.environ[\"DSP_NOTEBOOK_CACHEDIR\"] = repo_clone_path"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BiS0I8Oc9HDC"
      },
      "source": [
        "Before we start, we install DSPy and all dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u8WPRleVRkln"
      },
      "outputs": [],
      "source": [
        "%load_ext autoreload\n",
        "%autoreload 2\n",
        "\n",
        "import sys\n",
        "import os\n",
        "import regex as re\n",
        "import json\n",
        "\n",
        "import pkg_resources # Install the package if it's not installed\n",
        "if not \"dspy-ai\" in {pkg.key for pkg in pkg_resources.working_set}:\n",
        "    !pip install git+https://github.com/stanfordnlp/dspy.git\n",
        "    !pip install openai~=0.28.1\n",
        "\n",
        "from rich import print\n",
        "import dspy\n",
        "\n",
        "from dspy.evaluate import Evaluate\n",
        "from dspy.datasets.gsm8k import GSM8K, gsm8k_metric\n",
        "from dspy.teleprompt import BootstrapFewShotWithRandomSearch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xJcHuvwZ-rUk"
      },
      "source": [
        "### Step 0. **Getting Started** (Run the collapsed cells)\n",
        "\n",
        "We'll start by importing our dataset GSM8K, a dataset containing 8.5K high quality linguistically diverse grade school math word problems created by human problem writers. We preshuffled the dataset, and divided it into three smaller sets - train set, dev set (validation set), and test set. For simplicity, we will be using the train set and dev set.\n",
        "\n",
        "If you would like to inspect the dataset and see how we setup DSPy to use OpenAI gpt3, expand this step.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZQDVRJ4wgFQx"
      },
      "outputs": [],
      "source": [
        "gms8k = GSM8K()\n",
        "trainset, devset = gms8k.train, gms8k.dev\n",
        "len(trainset), len(devset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cYMPUeeGgg9_"
      },
      "source": [
        "Now we can inspect some examples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x3XWN_wfhgGW"
      },
      "outputs": [],
      "source": [
        "math_problem_example = devset[10]\n",
        "\n",
        "print(f\"Question: {math_problem_example.question}\\n\")\n",
        "print(f\"Gold Reasoning: {math_problem_example.gold_reasoning}\\n\")\n",
        "print(f\"Answer: {math_problem_example.answer}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "js5bCUuYcAks"
      },
      "source": [
        "Then we set up the language model (LM). **DSPy** supports multiple API and local models. In this notebook, we'll work with GPT-3.5 (`gpt-3.5-turbo`).\n",
        "\n",
        "We configure **DSPy** to use the turbo LM (`gpt-3.5-turbo`) by default. This can be overwritten for local parts of programs if needed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DBQg4TQBS7GH"
      },
      "outputs": [],
      "source": [
        "turbo = dspy.OpenAI(model='gpt-3.5-turbo', max_tokens=500)\n",
        "dspy.settings.configure(lm=turbo)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bCTCd0l1kQda"
      },
      "source": [
        "### Step 1. **First DSPy program**\n",
        "\n",
        "In **DSPy**, we will maintain a clean separation between **defining your modules in a declarative way** and **calling them in a pipeline to solve the task**.\n",
        "\n",
        "This allows you to focus on the information flow of your pipeline. **DSPy** will then take your program and automatically optimize **how to prompt** (or finetune) LMs **for your particular pipeline** so it works well.\n",
        "\n",
        "If you have experience with PyTorch, you can think of DSPy as the PyTorch of the foundation model space. Before we see this in action, let's first understand some key pieces."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p8Mth2j8kgVS"
      },
      "source": [
        "##### Using the Language Model: **Signatures** & **Predictors**\n",
        "\n",
        "Every call to the LM in a **DSPy** program needs to have a **Signature**.\n",
        "\n",
        "A signature consists of three simple elements:\n",
        "\n",
        "- A minimal description of the sub-task the LM is supposed to solve.\n",
        "- A description of one or more input fields (e.g., input question) that we will give to the LM.\n",
        "- A description of one or more output fields (e.g., the question's answer) that we will expect from the LM.\n",
        "\n",
        "Let's define a simple signature for basic math problem solving."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ue7--j2Tk5qe"
      },
      "outputs": [],
      "source": [
        "class SimpleMathSignature(dspy.Signature):\n",
        "    \"\"\"Answer the math question.\"\"\"\n",
        "\n",
        "    question = dspy.InputField(desc=\"A simple math question.\")\n",
        "    answer = dspy.OutputField(desc=\"The answer to the math question.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P-lX8Grsl9wT"
      },
      "source": [
        "In `SimpleMathSignature`, the docstring describes the sub-task here (i.e., answering math questions). Each `InputField` or `OutputField` can optionally contain a description `desc` too. When it's not given, it's inferred from the field's name (e.g., `question`).\n",
        "\n",
        "Notice that there isn't anything special about this signature in **DSPy**. We can just as easily define a signature that takes a long snippet from a PDF and outputs structured information, for instance.\n",
        "\n",
        "One trick for DSPy signature is that when it only contains simple fields performing straightforward tasks, we can replace the whole class definition with a syntactic sugar `question -> answer`. Now, lets define our first DSPy program with DSPy predictor. A predictor is a module that knows how to use the LM to implement a signature. Importantly, predictors can learn to fit their behavior to the task!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D3A1CH3foAFv"
      },
      "outputs": [],
      "source": [
        "basic_math_solver = dspy.Predict(\"question -> answer\") # Alternatively, we can write dspy.Predict(SimpleMathSignature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_GBP5jBUpRQP"
      },
      "source": [
        "**`DSPy.Predict`** is the simplest DSPy predictor. Now, we can call this minimal _program_ with a hand crafted question:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w1rZSbu2pQY3"
      },
      "outputs": [],
      "source": [
        "prediction = basic_math_solver(question=\"What is 1+1+1?\")\n",
        "\n",
        "print(f\"Answer: {prediction.answer}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WxFlRkxVqPw_"
      },
      "source": [
        "In the example above, we asked the predictor a simple math question \"What is 1+1+1?\". The model outputs an answer (\"3\").\n",
        "\n",
        "For visibility, we can inspect how this extremely basic predictor implemented our signature. Let's inspect the history of our LM (**turbo**)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7qMBWy-8qO0t"
      },
      "outputs": [],
      "source": [
        "_ = turbo.inspect_history(n=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VllCRN6tq-ty"
      },
      "source": [
        "Great. Now let's define the actual program. This is a class that inherits from `dspy.Module`.\n",
        "\n",
        "It needs two methods:\n",
        "\n",
        "- The `__init__` method will simply declare the sub-modules it needs: This time, we will be using a fancier predictor that implementes Chain-of-Thought prompting  `dspy.ChainOfThought`. `dspy.ChainOfThought` will add another field called \"rationale\" as output to help the model think step-by-step.\n",
        "- The `forward` method will describe the control flow of answering the question using the modules we have (here, we just have one)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kpxVuUblVvP7"
      },
      "outputs": [],
      "source": [
        "class SimpleMathSolver(dspy.Module):\n",
        "    def __init__(self):\n",
        "        self.prog = dspy.ChainOfThought(\"question -> answer\")\n",
        "\n",
        "    def forward(self, question):\n",
        "        pred = self.prog(question=question)\n",
        "        return pred\n",
        "\n",
        "simple_math_solver = SimpleMathSolver()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sUmsMEGtrjm7"
      },
      "source": [
        "#### **Exercise**\n",
        "Create your own math problem, use the math solver we just defined. Then, inspect the trace of the LM with `turbo.inspect_history` to see what has changed compared to the `dspy.Predict` predictor."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cXUH1WaOsIEr"
      },
      "outputs": [],
      "source": [
        "### Fill this code cell"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pVC0KPuvsoMD"
      },
      "source": [
        "We can now evaluate our simple math solver on the validation set.\n",
        "\n",
        "For a start, let's evaluate the accuracy of the predicted answer. We provide a simple metric function called `gsm8k_metric`, which essentially extract the numerical answer from the model input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EBOdvzCGTBYD"
      },
      "outputs": [],
      "source": [
        "evaluate = Evaluate(devset=devset[:], metric=gsm8k_metric, num_threads=16, display_progress=True, display_table=5)\n",
        "\n",
        "evaluate(simple_math_solver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CNFkEVkvtaWz"
      },
      "source": [
        "### Step 2. **Adding constraints with DSPy Assertions**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rFKOzC4NxkSe"
      },
      "source": [
        "We have **61.67%** on our validation set, not bad! But we also noticed two things: 1). many answers are sentences rather than the numerical result we want. Although we are able to parse most of the answers within `gsm8k_metric`, generating irrevalent tokens as answers might negatively affect the overall accuracy; 2). some of the reasoning might not contain the desired computational steps as in the example below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EMrMDwRXAdqi"
      },
      "outputs": [],
      "source": [
        "simple_math_solver(devset[0].question)\n",
        "_ = turbo.inspect_history()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "90BxOnIrAgnD"
      },
      "source": [
        "Forturnately, in DSPy, we can utilize a simple yet powerful construct called **LM Assertions** to constrain the output of LMs. For example, here, we can say:\n",
        "\n",
        "```python\n",
        "dspy.Suggest(len(pred.answer) < 10, \"Your Answer should be a number.\")\n",
        "```\n",
        "\n",
        "This suggestion tells the DSPy runtime that we expect the answer of our math solver to be short, and if the LM failed to yield such an answer, we instruct the LM that \"Your Answer should be a number.\"\n",
        "\n",
        "LM assertions in DSPy could either be a hard constraint `Assert` or a soft constraint `Suggest`. LM assertions accept two argument, one is the predicate to be tested, similar to that of traditional assertions; then, we also require an additional \"error message\" to guide the language model to refine itself when failing.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W2dBQ-g3YvO9"
      },
      "source": [
        "#### Math Solver with Suggestions\n",
        "We can encode the two observations we have into two suggestions, and add them to the `SimpleMathSolver`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UDCY7Jw3ys86"
      },
      "outputs": [],
      "source": [
        "def extract_number(question):\n",
        "    numbers = [int(s) for s in question.split() if s.isdigit()]\n",
        "    return numbers\n",
        "\n",
        "def has_numbers(rationale, numbers):\n",
        "    for number in numbers:\n",
        "        if str(number) not in rationale:\n",
        "            return False, number\n",
        "    return True, None\n",
        "\n",
        "class SimpleMathSolverWithSuggest(dspy.Module):\n",
        "    def __init__(self):\n",
        "        self.prog = dspy.ChainOfThought(\"question -> answer\")\n",
        "\n",
        "    def forward(self, question):\n",
        "        pred = self.prog(question=question)\n",
        "        rationale_has_numbers, missing_number = has_numbers(pred.rationale, extract_number(question))\n",
        "        dspy.Suggest(rationale_has_numbers, f\"Your Reasoning should contain {missing_number}.\")\n",
        "        dspy.Suggest(len(pred.answer) < 10, \"Your Answer should be a number.\")\n",
        "        return pred\n",
        "\n",
        "simple_math_solver_suggest = SimpleMathSolverWithSuggest().activate_assertions()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VWuOTcUIZLtj"
      },
      "source": [
        "Now we can rerun our math solver on the first question, and see how LM assertions in DSPy internally fix these errors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bANEA2jmFXpM"
      },
      "outputs": [],
      "source": [
        "simple_math_solver_suggest(devset[0].question)\n",
        "_ = turbo.inspect_history(n=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jzHHqOA0Zb9x"
      },
      "source": [
        "Finally, let's evaluate the performance of `simple_math_solver_suggest`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ASzhftKkZbHx"
      },
      "outputs": [],
      "source": [
        "evaluate(simple_math_solver_suggest)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "35xoJc9lye7x"
      },
      "source": [
        "### Step 3. **Compiling DSPy programs with optimizers**\n",
        "\n",
        "Another cool thing to do with DSPy is optimizers!\n",
        "\n",
        "A DSPy optimizer is an algorithm that can tune the parameters of a DSPy program (i.e., the prompts and/or the LM weights) to maximize the metrics you specify, like accuracy.\n",
        "\n",
        "There are many built-in optimizers in DSPy, which apply vastly different strategies. A typical DSPy optimizer takes three things:\n",
        "\n",
        "1. Your **DSPy program**. This may be a single module (e.g., dspy.Predict) or a complex multi-module program.\n",
        "\n",
        "2. Your **metric**. This is a function that evaluates the output of your program, and assigns it a score (higher is better).\n",
        "\n",
        "3. A few **training inputs**. This may be very small (i.e., only 5 or 10 examples) and incomplete (only inputs to your program, without any labels).\n",
        "\n",
        "If you happen to have a lot of data, DSPy can leverage that. But you can start small and get strong results.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FV5worzCvkyA"
      },
      "source": [
        "In this turtorial, we demonstrate one DSPy optimizer called `BootstrapFewShotWithRandomSearch`, which bootstraps demonstrations from the training set and search for the best combination of demonstrations. Two things to note here:\n",
        "1. Most optimizers work with LM assertions.\n",
        "2. This step is time/compute intensive. Therefore we cached the API calls. The good thing is, once you optmized the program, you can save the compiled DSPy program and reuse it later!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xF4Y2LMxnR1j"
      },
      "outputs": [],
      "source": [
        "optimizer = BootstrapFewShotWithRandomSearch(gsm8k_metric, max_bootstrapped_demos=3, max_labeled_demos=6, num_candidate_programs=6)\n",
        "\n",
        "compiled_prog = optimizer.compile(student=simple_math_solver, trainset=trainset[:], valset=devset[:100])\n",
        "compiled_prog_suggest = optimizer.compile(student=simple_math_solver_suggest, trainset=trainset[:], valset=devset[:100])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xiKwKBPJnEN6"
      },
      "outputs": [],
      "source": [
        "# Evaluating compiled program\n",
        "evaluate(compiled_prog)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nRRDcsV-AF8_"
      },
      "outputs": [],
      "source": [
        "# Evaluating compiled program with suggestions\n",
        "evaluate(compiled_prog_suggest)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mz4u8j-k2lNY"
      },
      "source": [
        "Now we inspect on our previous example, and see how the optmizer tunes the prompt:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Esd1Ry92vIt"
      },
      "outputs": [],
      "source": [
        "compiled_prog(devset[0].question)\n",
        "_ = turbo.inspect_history()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGG4S1TpUnFz"
      },
      "source": [
        "### More DSPy turtorials\n",
        "\n",
        "1. [Intro to DSPy](https://github.com/stanfordnlp/dspy/blob/main/intro.ipynb)\n",
        "2. [DSPy Assertions](https://colab.research.google.com/github/stanfordnlp/dspy/blob/main/examples/longformqa/longformqa_assertions.ipynb)\n",
        "3. [Quiz Generation](https://github.com/stanfordnlp/dspy/blob/main/examples/quiz/quiz_assertions.ipynb)\n",
        "4. ... more on [DSPy github](https://github.com/stanfordnlp/dspy)\n",
        "\n",
        "#### Contact: Shangyin Tan (shangyin@berkeley.edu)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "ghG4tCe3e-kQ",
        "xJcHuvwZ-rUk"
      ],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
